"""
PDF导出模块

用于将处理后的身份证图像导出为PDF格式
支持单页PDF和多页PDF合并功能
"""

import os
import numpy as np
from PIL import Image
from reportlab.lib.pagesizes import A4
from reportlab.platypus import SimpleDocTemplate, Spacer
from reportlab.platypus.flowables import Image as RLImage
from reportlab.lib.units import mm
from reportlab.lib.utils import ImageReader
import io


class PDFExporter:
    """
    PDF导出处理类
    
    功能：
    - 将图像转换为PDF格式
    - 支持A4页面布局
    - 保持图像质量和比例
    - 支持批量导出
    """
    
    def __init__(self):
        """初始化PDF导出器"""
        # A4纸尺寸 (mm)
        self.page_width = 210
        self.page_height = 297
        
        # PDF页面边距 (mm)
        self.margin_top = 20
        self.margin_bottom = 20
        self.margin_left = 15
        self.margin_right = 15
        
        # 可用内容区域
        self.content_width = self.page_width - self.margin_left - self.margin_right
        self.content_height = self.page_height - self.margin_top - self.margin_bottom
    
    def image_to_pdf(self, image_path: str, output_path: str, 
                     title: str = "身份证扫描件") -> bool:
        """
        将单张图像转换为PDF
        
        参数:
            image_path: 输入图像路径
            output_path: 输出PDF路径
            title: PDF文档标题
            
        返回:
            bool: 转换是否成功
            
        使用场景:
            exporter = PDFExporter()
            success = exporter.image_to_pdf("processed.jpg", "output.pdf")
        """
        try:
            # 创建PDF文档
            doc = SimpleDocTemplate(
                output_path,
                pagesize=A4,
                topMargin=self.margin_top * mm,
                bottomMargin=self.margin_bottom * mm,
                leftMargin=self.margin_left * mm,
                rightMargin=self.margin_right * mm,
                title=title
            )
            
            # 构建文档内容
            story = []
            
            # 添加图像
            img_element = self._create_image_element(image_path)
            if img_element:
                story.append(img_element)
            else:
                print(f"错误: 无法处理图像 {image_path}")
                return False
            
            # 生成PDF
            doc.build(story)
            
            print(f"PDF导出成功: {output_path}")
            return True
            
        except Exception as e:
            print(f"PDF导出失败: {e}")
            return False
    
    def numpy_to_pdf(self, image_array: np.ndarray, output_path: str,
                     title: str = "身份证扫描件") -> bool:
        """
        将numpy数组图像转换为PDF
        
        参数:
            image_array: 图像数组
            output_path: 输出PDF路径
            title: PDF文档标题
            
        返回:
            bool: 转换是否成功
            
        使用场景:
            exporter = PDFExporter()
            success = exporter.numpy_to_pdf(processed_array, "output.pdf")
        """
        try:
            # 将numpy数组转换为PIL图像
            if image_array.dtype != np.uint8:
                image_array = (image_array * 255).astype(np.uint8)
            
            pil_image = Image.fromarray(image_array)
            
            # 保存为临时文件然后转PDF
            temp_path = "temp_for_pdf.jpg"
            pil_image.save(temp_path, "JPEG", quality=95)
            
            # 转换为PDF
            success = self.image_to_pdf(temp_path, output_path, title)
            
            # 清理临时文件
            if os.path.exists(temp_path):
                os.remove(temp_path)
            
            return success
            
        except Exception as e:
            print(f"numpy数组转PDF失败: {e}")
            return False
    
    def batch_to_pdf(self, image_list: list, output_path: str,
                     title: str = "身份证批量扫描件") -> bool:
        """
        将多张图像合并为一个PDF文件
        
        参数:
            image_list: 图像路径列表
            output_path: 输出PDF路径
            title: PDF文档标题
            
        返回:
            bool: 转换是否成功
            
        使用场景:
            images = ["id1.jpg", "id2.jpg", "id3.jpg"]
            success = exporter.batch_to_pdf(images, "batch.pdf")
        """
        try:
            if not image_list:
                print("错误: 图像列表为空")
                return False
            
            # 创建PDF文档
            doc = SimpleDocTemplate(
                output_path,
                pagesize=A4,
                topMargin=self.margin_top * mm,
                bottomMargin=self.margin_bottom * mm,
                leftMargin=self.margin_left * mm,
                rightMargin=self.margin_right * mm,
                title=title
            )
            
            story = []
            
            # 添加每张图像
            for i, image_path in enumerate(image_list):
                if not os.path.exists(image_path):
                    print(f"警告: 图像文件不存在 - {image_path}")
                    continue
                
                # 添加图像
                img_element = self._create_image_element(image_path)
                if img_element:
                    story.append(img_element)
                    
                    # 如果不是最后一张图像，添加分页符
                    if i < len(image_list) - 1:
                        story.append(Spacer(1, self.content_height * mm))
                else:
                    print(f"警告: 无法处理图像 {image_path}")
            
            if not story:
                print("错误: 没有有效的图像可以添加到PDF")
                return False
            
            # 生成PDF
            doc.build(story)
            
            print(f"批量PDF导出成功: {output_path}")
            print(f"共处理 {len([s for s in story if isinstance(s, RLImage)])} 张图像")
            return True
            
        except Exception as e:
            print(f"批量PDF导出失败: {e}")
            return False
    
    def _create_image_element(self, image_path: str):
        """
        创建ReportLab图像元素
        
        参数:
            image_path: 图像文件路径
            
        返回:
            RLImage: ReportLab图像对象
        """
        try:
            # 打开图像获取尺寸
            with Image.open(image_path) as img:
                img_width, img_height = img.size
            
            # 计算缩放比例以适应页面
            scale_w = (self.content_width * mm) / img_width
            scale_h = (self.content_height * mm) / img_height
            scale = min(scale_w, scale_h, 1.0)  # 不放大
            
            # 计算最终显示尺寸
            display_width = img_width * scale
            display_height = img_height * scale
            
            # 创建ReportLab图像对象
            img_element = RLImage(
                image_path,
                width=display_width,
                height=display_height
            )
            
            return img_element
            
        except Exception as e:
            print(f"创建图像元素失败: {e}")
            return None
    
    def get_pdf_info(self, pdf_path: str) -> dict:
        """
        获取PDF文件信息
        
        参数:
            pdf_path: PDF文件路径
            
        返回:
            dict: PDF信息字典
            
        使用场景:
            info = exporter.get_pdf_info("output.pdf")
            print(f"文件大小: {info['file_size']}")
        """
        try:
            if not os.path.exists(pdf_path):
                return {"error": "PDF文件不存在"}
            
            file_size = os.path.getsize(pdf_path)
            
            return {
                "file_path": pdf_path,
                "file_size": f"{file_size / 1024:.1f} KB",
                "page_size": f"{self.page_width}mm x {self.page_height}mm",
                "margins": f"上下:{self.margin_top}mm, 左右:{self.margin_left}mm",
                "content_area": f"{self.content_width}mm x {self.content_height}mm"
            }
            
        except Exception as e:
            return {"error": f"获取PDF信息失败: {e}"}
    
    def create_combined_pdf(self, image_path: str, pdf_path: str,
                           include_original: bool = False) -> bool:
        """
        创建包含原图和处理结果的对比PDF
        
        参数:
            image_path: 处理后的图像路径
            pdf_path: 输出PDF路径
            include_original: 是否包含原图对比
            
        返回:
            bool: 创建是否成功
            
        使用场景:
            success = exporter.create_combined_pdf("processed.jpg", "report.pdf", True)
        """
        try:
            # 这个功能可以在后续版本中实现
            # 目前先实现基本的单图转PDF
            return self.image_to_pdf(image_path, pdf_path, "身份证处理报告")
            
        except Exception as e:
            print(f"创建对比PDF失败: {e}")
            return False 